import numpy as np
import matplotlib.pyplot as plt
import MDSplus


# 1. 下载数据
shot = 4417
tree = MDSplus.Tree('exl50u', 9381, path='192.168.20.11::/media/ennfusion/trees/exl50u')
tree.setTimeContext(-1, 3, 1e-5)
y_original = tree.getNode(r'\TEMP01').data()
y_original_time = tree.getNode(r'\TEMP01').dim_of().data()
y_true = tree.getNode(r'\HCN_NE001').data()
y_true_time = tree.getNode(r'\HCN_NE001').dim_of().data()
# Val1 = tree.getNode(r'\TEMP01').data()
# Time1 = tree.getNode(r'\TEMP01').dim_of().data()

# 2. 预处理
HCN = Val1.copy()
num = len(HCN)
time = Time1.copy()

# 3. 信号偏移
threshold1 = 1.3  # 噪声阈值
tag = 8  # 信号振幅
o_density = Val1 + 4.8  # 偏移信号
l_density = np.zeros(num)  # 初始化校正信号
l_density[0] = o_density[0]

# 4. 信号修正
for i in range(num - 1):
    tFlag = 1
    if abs(o_density[i + 1] - l_density[i]) <= threshold1:
        l_density[i + 1] = o_density[i + 1]
        tFlag = 0
    
    for k in range(1, 16):  # 处理周期性跳变
        if abs(o_density[i + 1] - tag * k - l_density[i]) <= threshold1:
            l_density[i + 1] = o_density[i + 1] - tag * k
            tFlag = 0
        if abs(o_density[i + 1] + tag * k - l_density[i]) <= threshold1:
            l_density[i + 1] = o_density[i + 1] + tag * k
            tFlag = 0
    
    if i > 20 and tFlag == 1:
        l_density[i + 1] = np.mean(l_density[i - 19:i + 1])  # 平滑处理

# 5. 归一化计算
dl = np.mean(l_density[:1000])
n_INT = (l_density - dl) / (4 * np.pi * 35.5 / 3.37 / 2)
nel = n_INT

# 6. 绘图
plt.figure()
plt.plot(Time1, nel)
plt.xlabel('Time (s)')
plt.ylabel('Density')
plt.title('Processed Density Signal')
plt.show()

print('Done!')
